package main

import (
	"bytes"
	_ "embed"
	"github.com/fogleman/gg"
	"github.com/nfnt/resize"
	ort "github.com/yalue/onnxruntime_go"
	"image"
	"image/color"
	"image/jpeg"
	_ "image/jpeg"
	"io"
	"os"
	"path"
	"runtime"
)

var (
	//go:embed libonnxruntime.so.1.16.3
	linuxLib []byte

	//go:embed onnxruntime.dll
	winLib []byte

	//go:embed bestplus20.onnx
	model1 []byte

	//go:embed bestplus50.onnx
	model2 []byte

	libName string
)

func init() {
	var lib []byte
	switch runtime.GOOS {
	case "linux":
		lib = linuxLib
		libName = "libonnxruntime.so.1.16.3"
	case "windows":
		lib = winLib
		libName = "onnxruntime.dll"
	default:
		panic("不支持的系统")
	}

	create, err := os.Create(libName)
	if err != nil {
		panic(err)
	}
	_, err = io.Copy(create, bytes.NewReader(lib))
	if err != nil {
		panic(err)
	}
	create.Close()

	file1, err := os.Create("bestplus20.onnx")
	if err != nil {
		panic(err)
	}
	defer file1.Close()

	file2, err := os.Create("bestplus50.onnx")
	if err != nil {
		panic(err)
	}
	defer file2.Close()

	io.Copy(file1, bytes.NewReader(model1))
	io.Copy(file2, bytes.NewReader(model2))
	str, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	err = InitEnv(path.Join(str, "onnxruntime.dll")) //拼文件路径
	if err != nil {
		panic(err)
	}
}

func InitEnv(libPath string) error {
	ort.SetSharedLibraryPath(libPath)
	err := ort.InitializeEnvironment()
	if err != nil {
		return err
	}
	return nil
}

type ProcessData struct {
	X    []float32
	Y    []float32
	W    []float32
	H    []float32
	Conf []float32
}

type Yolo struct {
}

func (l Yolo) ProcessImage(imgPath string, modelPath string) ([]byte, error) {
	file, err := os.ReadFile(imgPath)
	if err != nil {
		return nil, err
	}

	realImg, _, err := image.Decode(bytes.NewReader(file))
	if err != nil {
		return nil, err
	}

	img := resize.Resize(640, 640, realImg, resize.Lanczos2)

	outData, err := l.yoloProcess(img, modelPath, 0.8)
	if err != nil {
		return nil, err
	}

	return l.drawImage(outData, realImg, img)
}

func (l Yolo) drawImage(data ProcessData, realImg, img image.Image) ([]byte, error) {
	w := data.W
	h := data.H
	x := data.X
	y := data.Y
	conf := data.Conf

	realX := float64(realImg.Bounds().Dx())
	realY := float64(realImg.Bounds().Dy())

	ggCtx := gg.NewContextForImage(realImg)
	for i := 0; i < len(w); i++ {
		if conf[i] < 0.3 {
			continue
		}
		leftX := float64(x[i] - w[i]/2)
		leftY := float64(y[i] - h[i]/2)
		width := float64(w[i])
		hight := float64(h[i])

		ggCtx.DrawRectangle(realX*(leftX/640), realY*(leftY/640), realX*(width/640), realY*(hight/640))
		ggCtx.SetColor(color.RGBA{R: 255, A: 255})
	}

	ggCtx.Stroke()

	buf := bytes.NewBuffer(nil)

	err := jpeg.Encode(buf, ggCtx.Image(), nil)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (l Yolo) yoloProcess(img image.Image, modelPath string, conf float64) (out ProcessData, err error) {

	shape := ort.NewShape(1, 3, 640, 640)
	data := [3][640][640]float32{}

	for a := 0; a < 640; a++ {
		for b := 0; b < 640; b++ {
			R, G, B, _ := img.At(a, b).RGBA()
			data[0][a][b] = float32(R>>8) / 255.0
			data[1][a][b] = float32(G>>8) / 255.0
			data[2][a][b] = float32(B>>8) / 255.0
		}
	}

	d := make([]float32, 0, 3*640*640)
	for _, a := range data {
		for _, b := range a {
			for _, c := range b {
				d = append(d, c)
			}
		}
	}

	tensor, err := ort.NewTensor(shape, d)
	if err != nil {
		return
	}

	outData := make([]float32, 1*5*8400)
	outDataT, err := ort.NewTensor(ort.NewShape(1, 5, 8400), outData)
	if err != nil {
		return
	}

	session, err := ort.NewSession(modelPath, []string{"images"}, []string{"output0"}, []*ort.Tensor[float32]{tensor}, []*ort.Tensor[float32]{outDataT})
	if err != nil {
		return
	}
	err = session.Run()
	if err != nil {
		return
	}

	y := outData[:8400]
	x := outData[8400 : 2*8400]
	w := outData[2*8400 : 3*8400]
	h := outData[3*8400 : 4*8400]
	confData := outData[4*8400 : 5*8400]

	for idx, c := range confData {
		if float64(c) < conf {
			continue
		}

		out.Y = append(out.Y, y[idx])
		out.X = append(out.X, x[idx])
		out.W = append(out.W, w[idx])
		out.H = append(out.H, h[idx])
		out.Conf = append(out.Conf, confData[idx])
	}

	return
}
